"""
ARIMA模型
test: mae; 1172.88818359375， mape: 3.175144910812378, rmse: 1496.1046142578125
"""
import os
import numpy as np
import pandas as pd
from statsmodels.tsa.api import (stattools, ARIMA, ARMA, SARIMAX)
import matplotlib.pyplot as plt
from pylab import mpl
from keras.metrics import (RootMeanSquaredError,
                           MeanAbsoluteError,
                           MeanAbsolutePercentageError)
from model_evalute import Evalute, check_noise


# 指定默认字体
mpl.rcParams['font.sans-serif'] = ['SimHei']
# 解决保存图像是负号'-'显示为方块的问题
mpl.rcParams['axes.unicode_minus'] = False


def check_diff(data):
    """
    检查序列是否平稳
    :param data:
    :return:
    """
    """
    adfuller(Augmented Dickey-Fuller)检验
    原假设是非平稳，拒绝后就平稳了
    """
    adfuller_value, p_value, _, _, level, _ = stattools.adfuller(x=data)
    print(p_value, adfuller_value, level)
    if p_value < 0.05 or adfuller_value < level.get('1%', None):
        print('平稳')
    else:
        print('需要差分')

    """
    KPSS(Kwiatkowski-Phillips-Schmidt-Shin)检验
    原假设平稳，拒绝后就不平稳了
    """
    kpss_sta, p_value, lags, cirit = stattools.kpss(x=data)
    print(p_value, kpss_sta, cirit)
    if p_value < 0.05 or kpss_sta < cirit.get('1%', None):
        print("需要差分")
    else:
        print("平稳")


def choose_parameter_by_ic(data, max_ar, max_ma):
    """
    通过信息准则定阶(AIC/BIC)来选择最优参数（AIC/BIC最小的值）
    :param data:
    :param max_ar:
    :param max_ma:
    :return:
    """
    xic = stattools.arma_order_select_ic(y=data,
                                         max_ar=max_ar,
                                         max_ma=max_ma,
                                         ic=['aic', 'bic'])
    print('aic: {}, bic: {}'.
          format(xic.get('aic_min_order'), xic.get('bic_min_order')))


class ArmaModel:
    """
    ARMA模型
    """
    def __init__(self, train, test):
        self.train = train
        self.test = test

    def get_diff_train(self, diff_times, diff_periods, func):
        """
        差分
        :param diff_times: 差分次数
        :param diff_periods: 差分step
        :param func:
        :return:
        """
        train_data = func(self.train) if func else self.train
        train_diff = train_data.diff(diff_periods)[diff_periods:]
        diff_arr = [train_diff]
        for _ in range(diff_times - 1):
            mid_diff = train_diff.diff(diff_periods)[diff_periods:]
            diff_arr.append(mid_diff)
            train_diff = mid_diff
        return train_diff, diff_arr

    @staticmethod
    def build_model(train, order):
        """
        ARMA建模
        :param train:
        :param order:
        :return:
        """
        return ARMA(endog=train, order=order).fit()

    @staticmethod
    def predict_regain(pred, train, diff_arr, diff_periods):
        """
        ARMA预测差分的恢复
        :param pred:
        :param train:
        :param diff_arr:
        :param diff_periods:
        :return:
        """
        pred_raw = [0 for _ in range(len(pred))]
        for i in range(len(diff_arr) - 1, -1, -1):
            if i == 0:
                if set(pred_raw) != {0}:
                    for idx, num in enumerate(pred_raw):
                        if idx < len(train):
                            pred_raw[idx] = num + train[-len(train) + idx]
                        else:
                            pred_raw[idx] = num + pred_raw[idx - diff_periods]
                else:
                    for idx, num in enumerate(pred):
                        if idx < len(train):
                            pred_raw[idx] = num + train[-len(train) + idx]
                        else:
                            pred_raw[idx] = num + pred_raw[-diff_periods + idx]
            else:
                if set(pred_raw) != {0}:
                    for idx, num in enumerate(pred_raw):
                        pred_raw[idx + 1] = \
                            num + diff_arr[i - 1][-len(diff_arr[i - 1]) + idx]
                else:
                    for idx, num in enumerate(pred):
                        pred_raw[idx] = \
                            num + diff_arr[i - 1][-len(diff_arr[i - 1]) + idx]
                pred_raw.insert(0, diff_arr[i - 1][-len(diff_arr[i - 1])])
        return pred_raw

    def predict(self, diff_times, diff_periods, order, func, unfunc):
        """
        预测
        :param diff_times:
        :param diff_periods:
        :param order:
        :param func:
        :param unfunc: func的反函数
        :return:
        """
        train_diff, diff_arr = self.get_diff_train(diff_times=diff_times,
                                                   diff_periods=diff_periods,
                                                   func=func)
        model = self.build_model(train=train_diff, order=order)
        print(model.summary())
        # 预测
        # 二阶差分的predict貌似异常
        train_pred = model.predict()
        test_pred, stderr, conf_int = model.forecast(steps=len(self.test),
                                                     alpha=0.05)
        ci_low = conf_int[:, 0]
        ci_high = conf_int[:, 1]

        # 恢复
        train_data = func(self.train) if func else self.train
        last_train = train_data[-1 * diff_periods * diff_times:].tolist()

        train_pred_arr = \
            self.predict_regain(pred=train_pred,
                                train=train_data[:diff_periods * diff_times].
                                tolist(),
                                diff_arr=diff_arr,
                                diff_periods=diff_periods)
        test_pred_arr = self.predict_regain(pred=test_pred,
                                            train=last_train,
                                            diff_arr=diff_arr,
                                            diff_periods=diff_periods)
        test_pred_low = self.predict_regain(pred=ci_low,
                                            train=last_train,
                                            diff_arr=diff_arr,
                                            diff_periods=diff_periods)
        test_pred_high = self.predict_regain(pred=ci_high,
                                             train=last_train,
                                             diff_arr=diff_arr,
                                             diff_periods=diff_periods)

        train_pred_series = \
            pd.Series(train_pred_arr,
                      index=self.train.index[-len(train_pred_arr):])
        if unfunc:
            test_pred_arr = [unfunc(num) for num in test_pred_arr]
            test_pred_low = [unfunc(num) for num in test_pred_low]
            test_pred_high = [unfunc(num) for num in test_pred_high]

        return \
            pd.Series(test_pred_arr[diff_times - 1:], index=self.test.index), \
            test_pred_low, test_pred_high, train_pred_series

    def main(self, diff_times, diff_periods, order, pic_save_path,
             func=None, unfunc=None):
        """
        主函数
        :param diff_times: 差分次数
        :param diff_periods: 差分距离
        :param order:
        :param pic_save_path: 图片存储位置
        :param func: 训练集的处理
        :param unfunc: func的反函数
        :return:
        """
        pred_series, pred_low, pred_high, train_pred = \
            self.predict(diff_times=diff_times,
                         diff_periods=diff_periods,
                         order=order,
                         func=func,
                         unfunc=unfunc)
        # 画图
        self.train.plot()
        self.test.plot()
        train_pred.plot()
        pred_series.plot()
        plt.legend(['训练集', '测试集', '拟合', '预测'])
        # plt.legend(['训练集', '测试集', '预测'])
        pic_name = r'月度差分ARMA模型.png'
        plt.savefig(os.path.join(pic_save_path, pic_name))
        plt.show()
        #  评估
        train_err = train_pred - self.train
        # 白噪音检验
        check_noise(err=train_err)

        evalute = Evalute(y_true=self.test, y_pred=pred_series,
                          low=pred_low, high=pred_high)
        mae, mse, aic, aicc, bic = evalute.evalute_error_index()
        accuracy_score = evalute.accuracy()
        print('\n测试结果的参数评估：\n残差评估：mae: {}; mse: {}\n准确率：{}'.
              format(mae, mse, accuracy_score))


def arma_model(endog, order, test):
    """
    arma模型
    :param endog:
    :param order:
    :param test: test data
    :return:
    """
    model = ARMA(endog=endog, order=order).fit()
    print(model.summary())
    pred, stderr, conf_int = model.forecast(steps=len(test))
    pred_series = pd.Series(pred, index=test.index)

    # 画图
    endog.plot()
    test.plot()
    pred_series.plot()
    plt.legend(['训练集', '测试集', '预测'])
    plt.show()

    #  评估
    evalute = Evalute(y_true=test, y_pred=pred_series)
    mae, mse, aic, aicc, bic = evalute.evalute_error_index()
    print('\n测试结果的参数评估：\n残差评估：mae: {}; mse: {}'.format(mae, mse))


def arima_model(endog, order, test, pic_save_path):
    """
    arima模型
    :param endog:
    :param order:
    :param test: forecast
    :param pic_save_path:
    :return:
    """
    model = ARIMA(endog=endog, order=order).fit()
    print(model.summary())
    # 预测
    # train_pred = model.predict()
    pred, stderr, conf_int = model.forecast(steps=len(test), alpha=0.05)
    pred_series = pd.Series(pred, index=test.index)
    ci_low = conf_int[:, 0]
    ci_high = conf_int[:, 1]
    # 画图
    endog.plot()
    test.plot()
    pred_series.plot()
    plt.legend(['训练集', '测试集', '预测'])
    pic_name = r'xx.png'
    plt.savefig(os.path.join(pic_save_path, pic_name))
    plt.show()

    #  评估
    evalute = Evalute(y_true=test, y_pred=pred_series, low=ci_low, high=ci_high)
    mae, mse, aic, aicc, bic = evalute.evalute_error_index()
    accuracy_score = evalute.accuracy()
    print('\n测试结果的参数评估：\n残差评估：mae: {}; mse: {}\n准确率：{}'.
          format(mae, mse, accuracy_score))
    # evalute.evalute_predict_err()


def arima_model_month(train, order, seasonal_order, test, pic_save_path):
    """
    arima季度模型
    :param train:
    :param order:
    :param seasonal_order:
    :param test: forecast
    :param pic_save_path:
    :return:
    """
    model = SARIMAX(train, trend='ct',
                    order=order,
                    seasonal_order=seasonal_order).fit()
    print(model.summary())
    # 预测
    # train_pred = model.predict()
    pred = model.forecast(steps=len(test))
    pred_series = pd.Series(pred, index=test.index)
    # 画图
    train.plot()
    test.plot()
    pred_series.plot()
    plt.legend(['训练集', '测试集', '预测'])
    pic_name = r'SARIMA模型.png'
    # plt.savefig(os.path.join(pic_save_path, pic_name))
    plt.show()

    # 预测评估
    # mae
    mae = MeanAbsoluteError()
    mae.update_state(y_true=test, y_pred=pred)
    test_mae = mae.result()

    # mape
    mape = MeanAbsolutePercentageError()
    mape.update_state(y_true=test, y_pred=pred)
    test_mape = mape.result()

    # rmse
    rmse = RootMeanSquaredError()
    rmse.update_state(y_true=test, y_pred=pred)
    test_rmse = rmse.result()

    print('test: mae; {}， mape: {}, rmse: {}'.
          format(test_mae, test_mape, test_rmse))

    #  评估
    evalute = Evalute(y_true=test, y_pred=pred_series)
    mae, mse, aic, aicc, bic = evalute.evalute_error_index()
    print('\n测试结果的参数评估：\n残差评估：mae: {}; mse: {}'.
          format(mae, mse))


def arma_arima_model(data, pic_save_path):
    """

    :param data:
    :param pic_save_path: 图片存储路径
    :return:
    """
    # 检查平稳性
    # data.plot()
    # plt.show()
    data_log = np.log(data)
    # data_log = data_log.diff(12)[12:]
    # data_log.plot()
    # plt.show()

    two_diff_data = data_log.diff(1)[1:]
    two_diff_data = two_diff_data.diff(1)[1:]
    # print(data.tolist())
    # data.plot()
    # plt.show()
    month_diff_data = data_log.diff(12)[12:]
    # data.plot()
    # plt.show()

    # 检查序列是否平稳
    # check_diff(data=two_diff_data)

    # 通过信息准则定阶(AIC/BIC)来选择最优参数（AIC/BIC最小的值）
    """
    # month aic: 7,2; bic: 1,0
    # 二次差分：aic:4,4; bic:1,1
    # month log aic: 1,0; bic 1,0
    # 二次差分 log aic:3,3; bic: 1,1
    choose_parameter_by_ic(data=month_diff_data.loc[: '2019-08'],
                           max_ar=7,
                           max_ma=7)
    choose_parameter_by_ic(data=two_diff_data.loc[: '2019-08'],
                           max_ar=7,
                           max_ma=7)
    """

    # arma = ArmaModel(train=data.loc[: '2019-08'], test=data.loc['2019-09':])
    # 二次差分 arma模型
    # arma.main(diff_times=2, diff_periods=1, order=(1, 1),
    #           pic_save_path=pic_save_path)
    # 月度差分 arma模型
    # arma.main(diff_times=1, diff_periods=12, order=(1, 0),
    #           pic_save_path=pic_save_path)

    # 二次差分非季节性 arima模型
    # arima_model(endog=data.loc[: '2019-08'],
    #             order=(1, 2, 1),
    #             test=data.loc['2019-09':])

    # 季节性 arima模型
    arima_model_month(train=data.loc[: '2019-08'],
                      order=(0, 1, 0),
                      seasonal_order=(0, 1, 1, 12),
                      test=data.loc['2019-09':],
                      pic_save_path=pic_save_path)
